import cv2
import numpy as np
import h5py
import os
import glob


def list_to_npy_trainval(list_path, npy_path_im, npy_path_la, resize):
    train_folder_path = '/media/dell/delldisk/dell/wxm/Data/KaggleDDD/train'
    f_list = open(list_path, 'r')
    img_arrs = []
    labels = []
    for line in f_list:
        sub_path, label = line.split(' ')
        img_path = train_folder_path + '/' + sub_path
        img_arr = cv2.imread(img_path)
        img_arr = cv2.cvtColor(img_arr, cv2.COLOR_RGB2BGR)
        img_arr = img_arr[80:400, 320:]
        img_arr = img_arr[:, 80:560]
        img_arr = cv2.resize(img_arr, (resize, resize), interpolation=cv2.INTER_AREA)
        img_arrs.append(img_arr)
        labels.append(label)
    img_arrs = np.asarray(img_arrs, np.uint8)
    labels = np.asarray(labels, np.uint8)
    np.save(npy_path_im, img_arrs)
    np.save(npy_path_la, labels)


mean_vec = np.array([103.939, 116.779, 123.68], dtype=np.float32)
reshaped_mean_vec = mean_vec.reshape(3, 1, 1)


def list_to_h5_trainval(list_path, h5_path, resize):
    train_folder_path = '/media/lab/labdisk/home/lab1314/xryu/data/DistractedDriverDetection/imgs/train'
    # train_folder_path = '/media/dell/delldisk/dell/wxm/Data/KaggleDDD/train'
    f_list = open(list_path, 'r')
    img_arrs = []
    labels = []
    for line in f_list:
        sub_path, label = line.split(' ')
        img_path = train_folder_path + '/' + sub_path
        img_arr = cv2.imread(img_path)
        # img_arr = cv2.cvtColor(img_arr, cv2.COLOR_RGB2BGR)
        img_arr = img_arr[80:400, 320:]
        img_arr = img_arr[:, 80:560]
        img_arr = cv2.resize(img_arr, (resize, resize), interpolation=cv2.INTER_AREA) / 256
        img_arrs.append(img_arr)
        labels.append(label)
    f = h5py.File(h5_path)
    img_arrs = np.asarray(img_arrs, np.float32)
    img_arrs = img_arrs.transpose((0, 3, 1, 2))
    labels = np.asarray(labels, np.uint8)
    labels = labels.reshape((len(labels), 1))
    print img_arrs.shape
    print labels.shape
    # img_arrs = img_arrs.transpose((0, 3, 1, 2))  # input shape must be [batch_size, H, W, C]
    img_arrs = img_arrs[:] - reshaped_mean_vec
    f.create_dataset('data', data=img_arrs)
    f.create_dataset('label', data=labels)
    f.close()


def crop_resize_img_from_list(list_path, resize, new_list_path):
    train_folder_path = '/media/dell/delldisk/dell/Rui/data/Kaggle/DistractedDriverDetection/imgs/train'
    # train_folder_path = '/media/dell/delldisk/dell/wxm/Data/KaggleDDD/train'
    f_list = open(list_path, 'r')
    img_arrs = []
    labels = []
    subsubpath = ''
    f = open(new_list_path, 'w')
    for line in f_list:
        sub_path, label = line.split(' ')
        img_path = train_folder_path + '/' + sub_path
        img_arr = cv2.imread(img_path)
        # img_arr = cv2.cvtColor(img_arr, cv2.COLOR_RGB2BGR)
        # img_arr = img_arr[80:400, 320:]
        img_arr = img_arr[:, 80:560]
        img_arr = cv2.resize(img_arr, (resize, resize), interpolation=cv2.INTER_AREA)
        img_name = sub_path.split('/')[-1]
        if 'val' in list_path:
            subsubpath = 'val_cropped'
        elif 'train' in list_path:
            subsubpath = 'train_cropped'
        save_path = os.path.join(train_folder_path, subsubpath, str(img_name))
        line_path = os.path.join(subsubpath, str(img_name))
        new_line = line_path + ' ' + str(label)
        f.writelines(new_line)
        cv2.imwrite(save_path, img_arr)
    f.close()
    print 'image saved'
    print 'list writed done'


def read_list():
    train_cropped_path = '/media/lab/labdisk/home/lab1314/xryu/data/DistractedDriverDetection/imgs/train/train_cropped'
    val_cropped_path = '/media/lab/labdisk/home/lab1314/xryu/data/DistractedDriverDetection/imgs/train/val_cropped'
    train_files = glob.glob(train_cropped_path)
    val_files = glob.glob(val_cropped_path)


crop_resize_img_from_list(
    list_path='/media/dell/delldisk/dell/Rui/data/Kaggle/DistractedDriverDetection/fine-tune-caffe/data_txt/driver_sorted/'
    'val_list.txt',
    resize=256,
    new_list_path='/media/dell/delldisk/dell/Rui/data/Kaggle/DistractedDriverDetection/fine-tune-caffe/data_txt/driver_sorted/'
    'val_list_cropped.txt'
)
crop_resize_img_from_list(
    list_path='/media/dell/delldisk/dell/Rui/data/Kaggle/DistractedDriverDetection/fine-tune-caffe/data_txt/driver_sorted/'
    'train_list.txt',
    resize=256,
    new_list_path='/media/dell/delldisk/dell/Rui/data/Kaggle/DistractedDriverDetection/fine-tune-caffe/data_txt/driver_sorted/'
    'train_list_cropped.txt'
)
"""
list_to_h5_trainval(
    list_path='/media/lab/labdisk/home/lab1314/xryu/data/DistractedDriverDetection/fine-tune-caffe/data_txt/driver_sorted/val_list.txt',
    h5_path='/media/lab/labdisk/home/lab1314/xryu/data/DistractedDriverDetection/fine-tune-caffe/data_h5/val_caffe_RGB_mean_sub.h5',
    resize=224)
list_to_h5_trainval(
    list_path='/media/lab/labdisk/home/lab1314/xryu/data/DistractedDriverDetection/fine-tune-caffe/data_txt/driver_sorted/train_list.txt',
    h5_path='/media/lab/labdisk/home/lab1314/xryu/data/DistractedDriverDetection/fine-tune-caffe/data_h5/train_caffe_RGB_mean_sub.h5',
    resize=224)
"""
